
using DataFrames, Plots, Statistics, JSON, Distributed
plotlyjs();


const numcores = 14

if nprocs() < numcores
    addprocs(numcores - nprocs())
end

@show nprocs();

@everywhere include(joinpath(dirname(pwd()),"jgrc/TuLiPa/src/TuLiPa.jl"))
@everywhere include("JulES/src/JulES.jl");

nprocs() = 14


datayear = 2025
scenarioyearstart = 1981
scenarioyearstop = 2011
horizonstart = 1981
numscen = 7

tnormal = FixedDataTwoTime(getisoyearstart(datayear),getisoyearstart(horizonstart))

phaseinoffsetdays = 2
phaseinoffset = Millisecond(Day(phaseinoffsetdays)) # phase in straight away from second stage scenarios
phaseindelta = Millisecond(Week(5)) # Phase in the second stage scenario over 5 weeks
phaseinsteps = 5 # Phase in second stage scenario in 5 steps
tphasein = PhaseinTwoTime(getisoyearstart(datayear),getisoyearstart(horizonstart), getisoyearstart(horizonstart), phaseinoffset, phaseindelta, phaseinsteps);


sti_themadata = "data_fra_thema"

themastructure = json_to_elements(sti_themadata, "dataset_thema.json")
themaseries = json_to_elements(sti_themadata, "tidsserier_thema.json")

elements = vcat(themaseries,themastructure)
addscenariotimeperiod!(elements, scenarioyearstart, scenarioyearstop);


# Set horizons for price prognosis models
# Long
longhorizonend = getisoyearstart(horizonstart + 5)
longhydroperiodduration = Millisecond(Week(6))
longpowerparts = 6

longrhsdata = StaticRHSAHData("Power", datayear, scenarioyearstart, scenarioyearstop)
longmethod = KMeansAHMethod()
longclusters = 4
longunitduration = Millisecond(Hour(6))

# Medium
medweeklength = 54
medhorizonend = getisoyearstart(horizonstart) + Week(medweeklength)
medhydroperiodduration = Millisecond(Day(7)); @assert medweeklength % Int(longhydroperiodduration.value/3600000/168) == 0
medpowerparts = 7

medrhsdata = DynamicRHSAHData("Power")
medmethod = KMeansAHMethod()
medclusters = 4
medunitduration = Millisecond(Hour(4))

# Short
shortweeklength = 1
shorthorizonend = getisoyearstart(horizonstart) + Week(shortweeklength)
shorthydroperiodduration = Millisecond(Day(1)); @assert medweeklength % shortweeklength == 0
shortpowerparts = 12

# Preallocate results, distributed on cores
longprobs = distribute([HiGHS_Prob() for i in 1:numscen])
medprobs = distribute([HiGHS_Prob() for i in 1:numscen], longprobs)
shortprobs = distribute([HiGHS_Prob() for i in 1:numscen], longprobs)

medprices = distribute([Dict() for i in 1:numscen], longprobs)
shortprices = distribute([Dict() for i in 1:numscen], longprobs)

medendvaluesobjs = distribute([EndValues() for i in 1:numscen], longprobs)
nonstoragestates = distribute([Dict{StateVariableInfo, Float64}() for i in 1:numscen], longprobs)

# Organise inputs and outputs
probs = (longprobs, medprobs, shortprobs)
longinput = (longhorizonend, longhydroperiodduration, longrhsdata, longmethod, longclusters, longunitduration, "long")
medinput = (medhorizonend, medhydroperiodduration, medrhsdata, medmethod, medclusters, medunitduration, "long")
shortinput = (shorthorizonend, shorthydroperiodduration, shortpowerparts, "short")
allinput = (numcores, elements, horizonstart, tnormal, tphasein, phaseinoffsetdays)
output = (medprices, shortprices, medendvaluesobjs, nonstoragestates)
  
# Initialize price prognosis models and run for first time step. Run scenarios in parallell
@time pl_prognosis_init!(probs, allinput, longinput, medinput, shortinput, output)

 43.805738 seconds (3.75 M allocations: 184.031 MiB, 0.20% gc time, 2.74% compilation time)

Task (done) @0x0000000098f73e10


scenario = 1
index = collect(medprices[scenario]["steprange"])
prices = medprices[scenario]["matrix"]
labels = [name for name in medprices[scenario]["names"]]
p = plot(index,prices,label=reshape(labels, 1, length(labels)),legend=:outertopright)

for scenario in 2:numscen
    prices = medprices[scenario]["matrix"]
    labels = [name for name in medprices[scenario]["names"]]
    plot!(p,index,prices,label=reshape(labels, 1, length(labels)),legend=:outertopright)
end
display(p)


scenario = 1
index = collect(shortprices[scenario]["steprange"])
prices = shortprices[scenario]["matrix"]
labels = [name for name in shortprices[scenario]["names"]]
p = plot(index,prices,label=reshape(labels, 1, length(labels)),legend=:outertopright)

for scenario in 2:numscen
    prices = shortprices[scenario]["matrix"]
    labels = [name for name in shortprices[scenario]["names"]]
    plot!(p, index,prices,label=reshape(labels, 1, length(labels)),legend=:outertopright)
end
display(p)


# Read dataset for stochastic and market clearing models
sti_themadata = "data_fra_thema"
tst = JSON.parsefile(joinpath(sti_themadata, "dataset_thema_excl_hydro_nose.json")) # European dataset without aggregated hydro in Norway and Sweden
themastructure = getelements(tst) 
tse = JSON.parsefile(joinpath(sti_themadata, "tidsserier_thema.json"))
themaseries = getelements(tse, sti_themadata)

dse = JSON.parsefile("data_fra_dynmodell/tidsserier_detd.json")
detdseries = getelements(dse);
dst = JSON.parsefile("data_fra_dynmodell/dataset_detd.json")
detdstructure = getelements(dst);

detailedelements = vcat(themaseries,themastructure,detdseries,detdstructure)
addscenariotimeperiod!(detailedelements, scenarioyearstart, scenarioyearstop);


detailedrescopl = JSON.parsefile("data_fra_dynmodell/magasin_elspot1.json")
for k in keys(detailedrescopl) # quick fix, detailed dataset has reservoirs for SE4, aggregated does not, TODO: Improve mapping
    if detailedrescopl[k] == "SE4"
        detailedrescopl[k] = "SE3"
    end
end

scenario = 1
endvaluescopl = Dict()
instance = [getinstancename(getid(obj)) for obj in medendvaluesobjs[scenario].objects]
endvalues = medendvaluesobjs[scenario].values
endvaluescopl = Dict(instance .=> endvalues)

for (k,v) in detailedrescopl
    endvaluescopl["Reservoir_" * k] = endvaluescopl["Reservoir_" * v * "_hydro_reservoir"]
end


# Cut parameters
maxcuts = 13 # preallocate fixed number of cuts, no cut selection
lb = -1e10 # lower bound of the future value in the first iteration
reltol = 0.0001 # relative tolerance

# Parameters for stochastic subsystem problems
shorttotaldays = shortweeklength*7
medtotaldays = (medweeklength-2)*7 # we reuse prices for two weeks, so have to be two weeks shorter than price prognosis problem
shortstartstorage = 50
medstartstorage = 65
storageinfo = (shortstartstorage, medstartstorage, endvaluescopl)
shortterminputs = (detailedelements, shorttotaldays, numscen, horizonstart, phaseinoffsetdays, shortprices, true)
medterminputs = (detailedelements, medtotaldays, numscen, horizonstart, phaseinoffsetdays, medprices, false)

ustoragesystemobjects = []
ushorts = []
# Make modelobjects for short-term subsystems
@time makemastersubobjects!(shortterminputs, ustoragesystemobjects, ushorts)
# Make modelobjects for medium-term subsystems
@time makemastersubobjects!(medterminputs, ustoragesystemobjects, ushorts)

# Distribute subsystems with inputs and outputs on different cores
storagesystemobjects, shorts = distribute_subsystems(ustoragesystemobjects, ushorts)
masters = distribute([HiGHS_Prob() for i in 1:length(storagesystemobjects)], storagesystemobjects)
subs = distribute([[] for i in 1:length(storagesystemobjects)], storagesystemobjects)
states = distribute([Dict{StateVariableInfo, Float64}() for i in 1:length(storagesystemobjects)], storagesystemobjects)
cuts = distribute([SimpleSingleCuts() for i in 1:length(storagesystemobjects)], storagesystemobjects)
storagesystems = distribute([Dict() for i in 1:length(storagesystemobjects)], storagesystemobjects)

# Initialize subsystem problems and run for first time step. Run subsystems in parallell
@time pl_stochastic_init!(numcores, storagesystemobjects, shorts, masters, subs, states, cuts, storageinfo, numscen, lb, maxcuts, reltol, tphasein)

 11.533659 seconds (23.26 M allocations: 1.066 GiB, 2.77% gc time, 74.50% compilation time)
  3.094866 seconds (11.98 M allocations: 485.123 MiB, 6.09% gc time, 2.63% compilation time)
 63.126281 seconds (893.38 k allocations: 49.954 MiB, 0.44% compilation time)

Task (done) @0x0000000078de6100


# Bring data to local core
cutslocal = convert(Vector{SimpleSingleCuts}, cuts)
nonstoragestateslocal = convert(Vector{Dict}, nonstoragestates)

# Initialize market clearing problem and run for first time step
clearing, clearingstorages, nonstoragestatesmean, clearingendvaluesdict = clearing_init(detailedelements, tnormal, phaseinoffsetdays, cutslocal, nonstoragestateslocal);


# Initialize start states for next time step, also mapping to Thema storages and max capacity in Thema-storages
startstates, enekvglobaldict = startstates_init(clearing, detailedrescopl, longprobs[1], tnormal)

# Collect prices and start states
price = Dict()
powerhorizonix = argmax(getnumperiods(h) for h in clearing.horizons)
getareaprices!(price, clearing, clearing.horizons[powerhorizonix], tnormal)
areanames = price["names"]
ix = collect(price["steprange"])
pricematrix = price["matrix"]

statenames = collect(keys(startstates))
statematrix = collect(values(startstates));


tnormal += Day(2)
tphasein = PhaseinTwoTime(getdatatime(tnormal), getscenariotime(tnormal), getscenariotime(tnormal), phaseinoffset, phaseindelta, phaseinsteps);


@time pl_prognosis!(numcores, longprobs, medprobs, shortprobs, medprices, shortprices, nonstoragestates, startstates, tnormal, tphasein, Millisecond(0))

  3.745796 seconds (133.46 k allocations: 7.489 MiB, 1.52% compilation time)

Task (done) @0x0000000097a4e2f0


medpriceslocal = convert(Vector{Dict}, medprices)
shortpriceslocal = convert(Vector{Dict}, shortprices)
medendvaluesdict = getmedendvaluesdict(medendvaluesobjs)

@time pl_stochastic!(numcores, masters, subs, states, cuts, startstates, medpriceslocal, shortpriceslocal, medendvaluesdict, shorts, numscen, reltol, tphasein, Millisecond(0))

 13.893554 seconds (540.69 k allocations: 29.703 MiB, 1.02% compilation time)

Task (done) @0x0000000098c1f650


cutslocal = convert(Vector{SimpleSingleCuts}, cuts)
nonstoragestateslocal = convert(Vector{Dict}, nonstoragestates)

@time startstates = clearing!(clearing, clearingstorages, startstates, cutslocal, clearingendvaluesdict, nonstoragestateslocal, nonstoragestatesmean, detailedrescopl, enekvglobaldict)

updateareaprices!(price, clearing, clearing.horizons[powerhorizonix], tnormal)
ix = vcat(ix, collect(price["steprange"]))
pricematrix = vcat(pricematrix, price["matrix"])
statematrix = hcat(statematrix, collect(values(startstates)));

  1.478443 seconds (1.33 M allocations: 43.751 MiB, 23.52% compilation time)


# Only update medterm and longterm every 14 days (med/long prices and med/long watervalues and cuts)
skipmed = Millisecond(0)
tnormal += Day(2)
tphasein = PhaseinTwoTime(getdatatime(tnormal), getscenariotime(tnormal), getscenariotime(tnormal), phaseinoffset, phaseindelta, phaseinsteps)

steps = 24
# steps = (getisoyearstart(scenarioyearstop) - getisoyearstart(scenarioyearstart)).value/Millisecond(Day(2)).value
totaltime = @elapsed for i in 1:steps
# totaltime = @elapsed while getscenariotime(tnormal) < getisoyearstart(scenarioyearstop)
    
    display(tnormal)
    # Deterministic long/mid/short - calculate scenarioprices
    @time pl_prognosis!(numcores, longprobs, medprobs, shortprobs, medprices, shortprices, nonstoragestates, startstates, tnormal, tphasein, skipmed)
    
    # Stochastic sub systems - calculate storage value
    if skipmed.value == 0
        medpriceslocal = convert(Vector{Dict}, medprices)
        medendvaluesdict = getmedendvaluesdict(medendvaluesobjs)
    end
    shortpriceslocal = convert(Vector{Dict}, shortprices)

    @time pl_stochastic!(numcores, masters, subs, states, cuts, startstates, medpriceslocal, shortpriceslocal, medendvaluesdict, shorts, numscen, reltol, tphasein, skipmed)
    
    # Market clearing
    cutslocal = convert(Vector{SimpleSingleCuts}, cuts)
    nonstoragestateslocal = convert(Vector{Dict}, nonstoragestates)

    @time startstates = clearing!(clearing, clearingstorages, startstates, cutslocal, clearingendvaluesdict, nonstoragestateslocal, nonstoragestatesmean, detailedrescopl, enekvglobaldict)
    
    updateareaprices!(price, clearing, clearing.horizons[powerhorizonix], tnormal)
    ix = vcat(ix, collect(price["steprange"]))
    pricematrix = vcat(pricematrix, price["matrix"])
    statematrix = hcat(statematrix, collect(values(startstates)))
    
    skipmed += Millisecond(Day(2))
    if skipmed > Millisecond(Day(14))
        skipmed = Millisecond(0)
    end
    
    tnormal += Day(2)
    tphasein = PhaseinTwoTime(getdatatime(tnormal), getscenariotime(tnormal), getscenariotime(tnormal), phaseinoffset, phaseindelta, phaseinsteps)
end

display(string("The simulation took: ", totaltime/60, " minutes"))
display(string("Time usage per timestep: ", totaltime/steps, " seconds"))

FixedDataTwoTime(DateTime("2024-12-30T00:00:00"), DateTime("1981-01-02T00:00:00"))

  3.250508 seconds (31.73 k allocations: 2.177 MiB)
 12.420668 seconds (73.22 k allocations: 4.185 MiB)
  1.003028 seconds (1.02 M allocations: 26.095 MiB)

FixedDataTwoTime(DateTime("2024-12-30T00:00:00"), DateTime("1981-01-04T00:00:00"))

  1.794800 seconds (31.71 k allocations: 2.175 MiB)
  0.442913 seconds (72.89 k allocations: 4.173 MiB)
  0.879865 seconds (1.02 M allocations: 26.095 MiB)

FixedDataTwoTime(DateTime("2024-12-30T00:00:00"), DateTime("1981-01-06T00:00:00"))

  1.549399 seconds (31.70 k allocations: 2.175 MiB)
  0.791747 seconds (72.90 k allocations: 4.174 MiB)
  0.960900 seconds (1.02 M allocations: 26.095 MiB)

FixedDataTwoTime(DateTime("2024-12-30T00:00:00"), DateTime("1981-01-08T00:00:00"))

  1.341306 seconds (31.79 k allocations: 2.183 MiB)
  0.425484 seconds (72.89 k allocations: 4.173 MiB)
  0.925187 seconds (1.02 M allocations: 26.095 MiB)

FixedDataTwoTime(DateTime("2024-12-30T00:00:00"), DateTime("1981-01-10T00:00:00"))

  1.198269 seconds (31.79 k allocations: 2.183 MiB)
  1.215532 seconds (72.91 k allocations: 4.175 MiB)
  0.952580 seconds (1.02 M allocations: 26.095 MiB)


plot(ix,pricematrix,label=reshape(areanames, 1, length(areanames)),size=(800,500),legend=:outertopright)


# data = Dict()
# data["areanames"] = areanames
# data["pricematrix"] = pricematrix
# data["priceindex"] = ix

# data["statenames"] = statenames
# data["statematrix"] = transpose(statematrix)
# data["stateindex"] = [getisoyearstart(horizonstart) + Day(2*(i-1)) for i in 1:first(size(transpose(statematrix)))]

# open("JulES_2030_81_10.json", "w") do f
#     write(f, JSON.json(data))
# end

Import packages¶

Prepare parallell processing - import code on all cores¶

Starting point and scenario settings¶

Price prognosis models uses the Thema dataset¶

NB!¶

Initialize and run price prognosis models¶

Plot medium term prices for the 7 scenarios¶

Plot short term prices for the 7 scenarios¶

Stochastic subsystem models and market clearing model uses more detailed data¶

Get water values for detailed hydro (DETD) from aggregated hydro (Thema)¶

Stochastic subsystem models with Benders decomposition¶

Market clearing problem¶

Move to the next time step¶

Run price prognosis models for new time step¶

Run stochastic sub systems for new time step¶

Run market clearing model for new time step¶

Simulate next time steps¶

Plot resulting prices¶

Save results¶